from multiprocessing.pool import ThreadPool
from typing import Any, Dict, Optional
import psutil
from numpy import random
from scipy.stats import ortho_group
import numpy


class BaseANN(object):
    """Base class/interface for Approximate Nearest Neighbors (ANN) algorithms used in benchmarking."""

    def __init__(self):
        self.seed_mat = None  # 初始化seed_mat为None
        self.seed_mat_num = 100
        self._tablename = "items"

    @property
    def tablename(self):
        return self._tablename

    @tablename.setter
    def tablename(self, value):
        self._tablename = value
        self._update_query()  # 当tablename变化时更新查询语句
    
    @property
    def index_name(self):
        """获取索引名称"""
        return f"{self.tablename}_idx"

    def _update_query(self):
        """根据当前的metric和tablename更新查询语句
        该方法必须由子类重写，用于在metric或tablename变化时更新查询语句
        """
        raise NotImplementedError("子类必须实现_update_query方法")

    def done(self) -> None:
        """Clean up BaseANN once it is finished being used."""
        pass


    def get_memory_usage(self) -> Optional[float]:
        """Returns the current memory usage of this ANN algorithm instance in kilobytes.

        Returns:
            float: The current memory usage in kilobytes (for backwards compatibility), or None if
                this information is not available.
        """

        return psutil.Process().memory_info().rss / 1024

    def get_table_usage(self):
        """Get the size of the table"""
        return 0
    
    def drop(self) -> None:
        """Drop existing table"""
        pass

    def copy(self, X: numpy.array) -> None:
        """Copy data to table"""
        pass

    def create_index(self) -> None:
        """Create index"""
        pass


    def has_same_data(self, dataset) -> bool:
        """检查是否可以重用现有数据
        
        该方法应由子类重写。子类应实现具体的数据比较逻辑，判断现有数据是否与待加载数据集相同。
        
        Args:
            dataset: 数据集对象
            
        Returns:
            bool: 如果数据集相同可以重用返回True，否则返回False
        """
        try:
            # 1. 检查行数是否匹配
            train_data = dataset["train"]
            total_rows = train_data.shape[0]
            db_row_count = self.get_row_count()
            if db_row_count != total_rows:
                print(f"数据行数不匹配: 数据库={db_row_count}, 训练集={total_rows}")
                return False
                
            # 2. 检查样本点
            # 选择多个检查点
            check_points = [0]  # 第一个点
            if total_rows > 1:
                check_points.append(total_rows - 1)  # 最后一个点
            if total_rows > 10:
                check_points.append(total_rows // 2)  # 中间点
            
            for point_id in check_points:
                vector = train_data[point_id]
                
                # 获取数据库中的向量 - 由子类实现
                db_values = self.get_vector_by_id(point_id)
                if db_values is None:
                    print(f"数据库中找不到ID={point_id}的记录")
                    return False
                
                # 判断长度是否一致
                if len(db_values) != len(vector):
                    print(f"ID={point_id}的向量长度不同: 数据库={len(db_values)}, 训练集={len(vector)}")
                    return False
                
                # 向量比较 - 根据度量空间选择比较方法
                is_similar = self.compare_vectors(vector, db_values)
                if not is_similar:
                    return False
            
            print("数据集校验通过，可以重用现有数据")
            return True
            
        except Exception as e:
            print(f"检查数据集时出错: {e}")
            return False

    def get_row_count(self) -> int:
        """获取表中的行数
        
        子类必须实现此方法，查询并返回数据表中的行数。
        
        Returns:
            int: 表中的行数
        """
        raise NotImplementedError("子类必须实现get_row_count方法")
    
    def get_vector_by_id(self, vector_id) -> numpy.ndarray:
        """根据ID获取向量
        
        子类必须实现此方法，从数据库获取指定ID的向量并返回为numpy数组。
        
        Args:
            vector_id: 向量的ID
            
        Returns:
            numpy.ndarray: 向量数据，如果找不到则返回None
        """
        raise NotImplementedError("子类必须实现get_vector_by_id方法")
    
    def compare_vectors(self, vec1, vec2) -> bool:
        """比较两个向量是否相似
        
        根据当前的度量空间(metric)选择合适的比较方法:
        - 对于angular空间，使用余弦相似度
        - 对于euclidean空间，使用L2距离
        
        Args:
            vec1: 第一个向量（numpy数组）
            vec2: 第二个向量（numpy数组）
            
        Returns:
            bool: 如果向量相似返回True，否则返回False
        """
        try:
            # 获取当前的度量空间类型
            metric = getattr(self, '_metric', 'euclidean')  # 默认为欧几里得空间
            
            if metric == "angular":
                # 使用余弦相似度比较
                norm1 = numpy.linalg.norm(vec1)
                norm2 = numpy.linalg.norm(vec2)
                if norm1 == 0 or norm2 == 0:
                    similarity = 1.0 if norm1 == norm2 else 0.0
                else:
                    similarity = numpy.dot(vec1, vec2) / (norm1 * norm2)
                
                # 余弦相似度接近1表示向量相似
                if similarity < 0.999:  # 设置一个合理的阈值
                    print(f"向量余弦相似度为{similarity}，低于阈值0.999，判定为不同向量")
                    return False
                else:
                    print(f"向量余弦相似度为{similarity}，判定为相同向量")
                    return True
            else:
                # 使用L2距离比较
                distance = numpy.sqrt(numpy.sum((vec1 - vec2) ** 2))
                
                # 距离接近0表示向量相似
                if distance > 0.1:  # 设置一个合理的阈值
                    print(f"向量L2距离为{distance}，超过阈值0.1，判定为不同向量")
                    return False
                else:
                    print(f"向量L2距离为{distance}，判定为相同向量")
                    return True
                
        except Exception as e:
            print(f"向量比较出错: {e}")
            return False

    def has_same_index(self) -> bool:
        """检查是否有相同参数的索引
        
        该方法应由子类重写。子类应实现具体的索引检查逻辑，判断现有索引是否有效且参数匹配。
        
        Returns:
            bool: 如果索引存在且有效返回True，否则返回False
        """
        return False

    def query(self, q: numpy.array, n: int) -> numpy.array:
        """Performs a query on the algorithm to find the nearest neighbors. 

        Note: This is a placeholder method to be implemented by subclasses.

        Args:
            q (numpy.array): The vector to find the nearest neighbors of.
            n (int): The number of nearest neighbors to return.

        Returns:
            numpy.array: An array of indices representing the nearest neighbors.
        """
        return []  # array of candidate indices

    def batch_query(self, X: numpy.array, n: int) -> None:
        """Performs multiple queries at once and lets the algorithm figure out how to handle it.

        The default implementation uses a ThreadPool to parallelize query processing.

        Args:
            X (numpy.array): An array of vectors to find the nearest neighbors of.
            n (int): The number of nearest neighbors to return for each query.
        Returns: 
            None: self.get_batch_results() is responsible for retrieving batch result
        """
        pool = ThreadPool()
        self.res = pool.map(lambda q: self.query(q, n), X)

    def get_batch_results(self) -> numpy.array:
        """Retrieves the results of a batch query (from .batch_query()).

        Returns:
            numpy.array: An array of nearest neighbor results for each query in the batch.
        """
        return self.res

    def get_additional(self) -> Dict[str, Any]:
        """Returns additional attributes to be stored with the result.

        Returns:
            dict: A dictionary of additional attributes.
        """
        return {}

    def __str__(self) -> str:
        return self.name

    def gen_vec(self, mag, dim):
        vec = numpy.random.randn(dim)
        vec /= numpy.linalg.norm(vec)
        return vec * mag

    def _transform(self, vec):
        if self.seed_mat is None:
            self.seed_mat = ortho_group.rvs(len(vec), self.seed_mat_num)
        
        random.seed(random.randint(1000000))

        # 生成一个0.01到0.1之间的随机幅度
        mag = random.uniform(0.01, 0.1)
                
        if random.rand() < 0.5:
            # 50%概率:用预生成的正交矩阵做线性变换,再加上原向量
            return self.seed_mat[random.randint(self.seed_mat_num)] @ vec * mag + vec
        else:
            # 50%概率:直接加上一个随机向量
            return vec + self.gen_vec(mag, len(vec))

    def transform(self, vec):
        # 先进行随机变换
        new_vec = self._transform(vec)
        # 再归一化为单位向量
        new_vec /= numpy.linalg.norm(new_vec)
        return new_vec